'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 7_study_parent_doc_retriever.py
* @Time: 2025/10/31
* @All Rights Reserve By Brtc
'''
import dotenv
import weaviate
from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import LocalFileStore
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey

dotenv.load_dotenv()

#1、创建文档加载器
loaders = [UnstructuredFileLoader("./eshop_goods.txt"),
           UnstructuredFileLoader("./01.项目API文档.md")]
docs = []
for loader in loaders:
    docs.extend(loader.load())
#2、创建文本分割器
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
#3、创建向量数据库与文档数据库
embedding = OpenAIEmbeddings(model="text-embedding-3-small")

client = weaviate.connect_to_weaviate_cloud(
    skip_init_checks=True,
    cluster_url="https://g2erxw7ety2jjvovjy2d0g.c0.asia-southeast1.gcp.weaviate.cloud",
    auth_credentials=AuthApiKey("TzRXakNkVEF1RnZJTWMvbl9QcERLY2lUbmRESzhjbjF4U1o4YWVwYnRHUTh2ME5FeUdGOFlhYjUzMUw4PV92MjAw")
)
vector_store  = WeaviateVectorStore(client=client, index_name="Test2Parent", text_key="text", embedding=embedding)
store = LocalFileStore("./parent-document")
#4、创建父类文档检索器
retriever = ParentDocumentRetriever(
    vectorstore=vector_store,
    byte_store=store,
    child_splitter=text_splitter,
    search_kwargs={"k":2}
)
#5、添加文档
#retriever.add_documents(docs, ids = None)

#6检索并返回文档
search_docs = retriever.invoke("请推荐一些潮州特产")
for doc in search_docs:
    print("===================================")
    print(doc.page_content[:50])
    print(doc.metadata)

